import torch
from env import JammingEnv
from model import QNet
from config import *
import numpy as np

def main():
    env = JammingEnv()
    # 加载训练好的模型
    agent1 = QNet(env.n_sub, env.n_sub).to(DEVICE)
    agent2 = QNet(env.n_sub + 1 + len(env.targets), BAND_OPTS).to(DEVICE)
    agent1.load_state_dict(torch.load("agent1.pth", map_location=DEVICE))
    agent2.load_state_dict(torch.load("agent2.pth", map_location=DEVICE))
    agent1.eval(); agent2.eval()

    state = env.reset()
    s1 = state["s1"]
    with torch.no_grad():
        while True:
            a1 = agent1(torch.tensor(s1, device=DEVICE).unsqueeze(0)).argmax().item()
            s2 = np.concatenate([s1, [a1], state["s2_base"]]).astype(np.float32)
            a2 = agent2(torch.tensor(s2, device=DEVICE).unsqueeze(0)).argmax().item()

            s1, _, done = env.step(a1, a2)
            s1 = s1["s1"]
            if done:
                break
    env.summary()

if __name__ == "__main__":
    main()